# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import datetime
import functools
import logging

from coin.base.timestamp import get_timestamp

from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.shared.kr_rest.product import SharedCurrency
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.strategy.hk_hard_arb.og_util import print_working_orders

from coin.strategy.cheese.config import L1Config
from coin.strategy.cheese.pricer import PricerBuilder
from coin.strategy.cheese.sizers.l1 import SpotL1Sizer, FuturesL1Sizer
from coin.strategy.cheese.logger import StrategyLoggerWrapper
from coin.strategy.mm.live_envs import get_live_env_cls

from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import create_live_env_instance
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.cheese.executor import L2Executor
from coin.support.accounting.logic.account_proto_util import gen_balance_with_reserve

from coin.strategy.cheese.pnl_calculator import PnlCalculator
from coin.exchange.okex.kr_rest.currency import OkexCurrency

from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls


class BaseStrategy(object):
  def __init__(self, strategy_params):
    self._logger = logging.getLogger(__name__)

    self._feed_available = False
    self._og_warn = True

    exchange_0 = strategy_params['exchange_0']
    exchange_1 = strategy_params['exchange_1']

    if exchange_0 == exchange_1:
      self._exchanges = [exchange_0]
    else:
      self._exchanges = [exchange_0, exchange_1]

    self._subreqs = [SubscriptionRequest.from_any(exchange) for exchange in self._exchanges]
    self._book_builder_clses = [get_flow_book_builder_cls(subreq) for subreq in self._subreqs]
    self._subreq_id_map = {subreq: req_id for req_id, subreq in enumerate(self._subreqs)}

    self._live_env_clses = None
    self._l2_executors = {}

    self._live_env = None
    self.og = None

    # Config
    config = strategy_params['config']
    self._og_params = config.get('og_params', {})

    # StrategyLogger
    name = strategy_params['name']
    revision = strategy_params['revision']
    self._strategy_logger = StrategyLoggerWrapper(name=name, revision=revision)

    # Pricer
    self._l1_config = L1Config(exchanges=self._exchanges,
                               config=config,
                               subreq0=self._subreqs[0],
                               subreq1=self._subreqs[-1])

    self._ready = False

    # Futures
    self.is_futures = exchange_0.startswith('Futures')

    self._l1_pricer = PricerBuilder(config).build(self._l1_config, self._strategy_logger)
    if self.is_futures:
      self._l1_sizer = FuturesL1Sizer(l1_config=self._l1_config)
    else:
      self._l1_sizer = SpotL1Sizer(l1_config=self._l1_config)

    feed_checker_config_map = {}
    if len(self._subreqs) == 1:
      products = []
      if self._l1_config.extra_product:
        products = products + [self._l1_config.extra_product]
      for exchange_id in (0, 1):
        products = products + self._l1_config.get_native_products(exchange_id)
      subscription_key = get_subscription_key(products[0])
      feed_checker_config_map[subscription_key] = get_default_feed_checker_config(products)
    else:
      for exchange_id in (0, 1):
        products = self._l1_config.get_native_products(exchange_id)
        subscription_key = get_subscription_key(products[0])
        if self._l1_config.extra_product and self._l1_config.extra_product.exchange_id == exchange_id:
          products = products + [self._l1_config.extra_product]
        feed_checker_config_map[subscription_key] = get_default_feed_checker_config(products)

    self._feed = MultiFeedSubsystem(feed_checker_config_map=feed_checker_config_map)
    self._live_env_cls = get_live_env_cls(self._subreqs[0])
    self._og_config = strategy_params['og_config']

    self._reserve_strat_balance = gen_balance_with_reserve(
        {key: l2_config.reserve_balance for key, l2_config in self._l1_config.items()})

    self._warning = True
    products = self._l1_config.get_native_products(0)
    self._bookmap = BookMap(products)
    self._balance_printer = get_pnl_balance_printer_cls(self._subreqs[0])(
        products, self._bookmap, self._l1_config.reserve_map())
    self._reserve_proto = self._l1_config.reserve_proto()

  def touch_account_request(self):
    self._logger.info('ACCOUNT REQ')
    self.og.og_logger.write_account_request()
    self._strategy_logger.write_request()
    self._strategy_logger.write_reserve(self._reserve_proto)
    self._ioloop.add_timeout(datetime.timedelta(hours=1), self.touch_account_request)

  def print_balance(self, periodic=False):
    timestamp = get_timestamp()
    self._logger.info('Pricer: %s, Feed: %s', self._l1_pricer.available(), self._feed.available)
    if self._l1_pricer.available() and self._feed.available:
      self._balance_printer.print_pnl_balance()
    if self.og:
      print_working_orders(self._logger, self.og)

    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=30), self.print_balance, True)

  def log_balance(self, periodic=False):
    pnl_snapshot_proto = None
    pnl_calculator = self._l1_pricer._pnl_calculator
    if pnl_calculator is not None:
      pnl_snapshot_proto = pnl_calculator.get_snapshot_proto()
    if pnl_snapshot_proto is not None:
      self._strategy_logger.write_pnl_snapshot(0, pnl_snapshot_proto)
    else:
      self._logger.warn('pnl_snapshot is None')

    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=3), self.log_balance, True)

  def subscribe_feeds(self):
    for key in self._l1_config.keys():
      self._l2_executors[key] = L2Executor(self._l1_config,
                                           self._l1_config.get_l2(key),
                                           l2_pricer=self._l1_pricer.get_l2(key),
                                           l2_sizer=self._l1_sizer.get_l2(key),
                                           strategy_logger=self._strategy_logger,
                                           balance_printer=self._balance_printer,
                                           logger=self._logger)

    for exchange_id in (0, 1):
      for product in self._l1_config.get_native_products(exchange_id):
        self._feed.subscribe_book(
            FeedSubscriptionRequest.create([product], self._subreqs[-exchange_id]),
            functools.partial(self._on_book, product, exchange_id))
        self._feed.subscribe_trade(
            FeedSubscriptionRequest.create([product], self._subreqs[-exchange_id]),
            functools.partial(self._on_trade, product, exchange_id))

    if self._l1_config.extra_product:
      exchange_id = self._l1_config.extra_product.exchange_id
      self._feed.subscribe_book(
          FeedSubscriptionRequest.create([self._l1_config.extra_product],
                                         self._subreqs[exchange_id]),
          functools.partial(self._on_book, self._l1_config.extra_product, exchange_id))

  def subscribe_live_env(self, ioloop, queue, flow_sub):
    sub_key = self._subreqs[0]
    og_params = convert_og_config_from_dict(sub_key, self._og_params)
    products = self._l1_config.get_native_products(0)
    currencies = self._l1_config.get_currencies(0)

    self._live_env = create_live_env_instance(self._live_env_cls,
                                              currencies=currencies,
                                              trade_products=products,
                                              og_trade_key=self._og_config,
                                              og_params=og_params,
                                              og_log_callback=functools.partial(
                                                  self._strategy_logger.write_og_log, 0),
                                              ioloop=ioloop,
                                              queue=queue,
                                              flow_sub=flow_sub)
    self._live_env.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_order_gateway(self, subreq, order_gateway):
    self._logger.info('OG initialized: %s', subreq)
    if subreq.exchange == 'Binance':  # hack
      subreq = SubscriptionRequest('Spot', 'Binance', 'snapshot')

    subreq_id = self._subreq_id_map[subreq]
    assert subreq_id == 0, subreq_id
    self.og = order_gateway
    if self.og is None:
      return

    self._balance_printer.register_order_gateway(self.og)

    keys = self._l1_config.keys()
    reserves = {key: self._l1_config.get_l2(key).reserve_balance for key in keys}
    # TODO(inkyu): No hope to generalize.
    bnb_reserve = self._l1_config.bnb_reserve
    base_currencies = self._l1_config.get_base_currencies()
    if bnb_reserve is not None:
      bnb_currency = SharedCurrency.FromStrCurrency('BNB')
      reserves[bnb_currency.currency] = bnb_reserve
      base_currencies = base_currencies + [bnb_currency]
      reserves['BNB'] = bnb_reserve
    base_currency_reserves = {  # TODO(inkyu): Remove
        SharedCurrency.FromStrCurrency(key): reserve for key,
        reserve in reserves.items()
    }

    extra_currencies = {}
    okdk_value = self._l1_config.okdk_value
    if okdk_value is not None:
      reserves['OKDK'] = 0

    # Pass quote currency only if we trade on spot product
    quote_currency = None
    if not self._l1_config.is_futures:
      quote_currency = self._l1_config.quote_currencies[0]
    self._l1_pricer.set_pnl_calculator(
        PnlCalculator(quote_currency=quote_currency,
                      reserves=reserves,
                      og=self.og,
                      balance_printer=self._balance_printer))
    if okdk_value is not None:
      self._l1_pricer._pnl_calculator.set_price('OKDK', okdk_value)
      extra_currencies = {OkexCurrency.FromStrCurrency('OKDK'): okdk_value}

    self._ready = True
    for sym, product_executor in self._l2_executors.items():
      product_executor.on_og_reset(self.og)

    self.touch_account_request()

  def _on_trade(self, product, exchange_id, trade):
    key = SharedCurrency.FromCurrency(product.base).currency

    feed_available = self._feed.available
    og_available = self.og is not None
    self._l1_pricer.on_trade(key, exchange_id, trade, feed_available and og_available)

  def _on_book(self, product, exchange_id, book):
    self._bookmap.set_book(product, book)

    feed_available = self._feed.available
    if feed_available != self._feed_available:
      if feed_available:
        self._logger.info('Strategy turned ON.')
      else:
        self._logger.info('Strategy turned OFF.')

    self._feed_available = self._feed.available

    if product == self._l1_config.extra_product:
      if exchange_id == 0 and self._l1_pricer._pnl_calculator is not None:
        self._l1_pricer._pnl_calculator.on_book(product, book)
      return

    og_available = True
    key = SharedCurrency.FromCurrency(product.base).currency
    if not self.og:
      if self._og_warn:
        self._logger.info('OG %s not ready' % self._exchanges[0])
        self._og_warn = False
      og_available = False
    else:
      self._l2_executors[key].process_position()

    self._l1_pricer.on_book(product, exchange_id, book, feed_available and og_available)
    self._l2_executors[key].on_book(product, exchange_id, book, feed_available and og_available)

    if feed_available and og_available:
      self._l2_executors[key].execute(book.timestamp)
      self._l1_pricer.update_quote_value(balance_printer=self._balance_printer)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env.start()
    self.print_balance(periodic=True)
    self.log_balance(periodic=True)
